package practice.graph;

import algorithm.unionfind.general.UnionFind;
import datastructure.heap.BinaryHeap;

import java.util.*;


/**
 * @Description:
 * @Author Ammar
 * @Create 2023/5/19 18:43
 */
public class ListGraph<V, E> extends Graph<V, E> {
    private Map<V, Vertex<V, E>> vertexs = new HashMap<>();
    private Set<Edge<V, E>> edges = new HashSet<>();

    private Comparator<Edge<V, E>> edgeComparator = (Edge<V, E> e1, Edge<V, E> e2) -> {
        return weightManager.compare(e1.weight, e2.weight);
    };

    public ListGraph() {
        this(null);
    }

    public ListGraph(WeightManager<E> weightManager) {
        super(weightManager);
    }

    @Override
    public int edgesSize() {
        return edges.size();
    }

    @Override
    public int verticesSize() {
        return vertexs.size();
    }

    @Override
    public void addVertex(V v) {
        if (vertexs.containsKey(v)) return;
        vertexs.put(v, new Vertex<>(v));
    }

    @Override
    public void addEdge(V from, V to) {
        addEdge(from, to, null);
    }

    @Override
    public void addEdge(V from, V to, E weight) {
        Vertex<V, E> fromV = vertexs.get(from);
        if (fromV == null) {
            fromV = new Vertex<>(from);
            vertexs.put(from, fromV);
        }
        Vertex<V, E> toV = vertexs.get(to);
        if (toV == null) {
            toV = new Vertex<>(to);
            vertexs.put(to, toV);
        }
        Edge<V, E> edge = new Edge<>(fromV, toV);
        edge.weight = weight;
        if (fromV.outEdges.remove(edge)) {
            toV.inEdges.remove(edge);
            edges.remove(edge);
        }
        fromV.outEdges.add(edge);
        toV.inEdges.add(edge);
        edges.add(edge);
    }

    @Override
    public void removeVertex(V v) {
        Vertex<V, E> remove = vertexs.remove(v);
        if (remove == null) return;
        Iterator<Edge<V, E>> outI = remove.outEdges.iterator();
        while (outI.hasNext()) {
            Edge<V, E> edge = outI.next();
            edge.to.inEdges.remove(edge);
            outI.remove();
            edges.remove(edge);
        }
        Iterator<Edge<V, E>> inI = remove.inEdges.iterator();
        while (inI.hasNext()) {
            Edge<V, E> edge = inI.next();
            edge.from.outEdges.remove(edge);
            inI.remove();
            edges.remove(edge);
        }
    }

    @Override
    public void removeEdge(V from, V to) {
        Edge<V, E> edge = new Edge<>(vertexs.get(from), vertexs.get(to));
        if (edges.remove(edge)) {
            vertexs.get(from).outEdges.remove(edge);
            vertexs.get(to).inEdges.remove(edge);
        }
    }

    @Override
    public void bfs(V begin, VertexVisitor<V> visitor) {
        Vertex<V, E> beginV = (Vertex<V, E>) vertexs.get(begin);
        if (beginV == null) return;
        Set<Vertex<V, E>> visitedVertex = new HashSet<>();
        Queue<Vertex<V, E>> queue = new LinkedList<>();
        queue.offer(beginV);
        visitedVertex.add(beginV);
        while (!queue.isEmpty()) {
            Vertex<V, E> vertex = queue.poll();
            if (visitor.visit(vertex.value)) return;
            for (Edge<V, E> outEdge : vertex.outEdges) {
                if (visitedVertex.contains(outEdge.to)) continue;
                queue.offer(outEdge.to);
                visitedVertex.add(outEdge.to);
            }
        }
    }

    @Override
    public void dfs(V begin, VertexVisitor<V> visitor) {
        Vertex<V, E> beginV = (Vertex<V, E>) vertexs.get(begin);
        if (beginV == null) return;
        dfs(beginV, new HashSet<>(), visitor);
    }

    private void dfs(Vertex<V, E> vertex, Set<Vertex<V, E>> visitedVertexs, VertexVisitor<V> visitor) {
        if (visitor.visit(vertex.value)) return;
        visitedVertexs.add(vertex);
        for (Edge<V, E> outEdge : vertex.outEdges) {
            if (visitedVertexs.contains(outEdge.to)) continue;
            dfs(outEdge.to, visitedVertexs, visitor);
        }
    }

    @Override
    public void dfs1(V begin, VertexVisitor<V> visitor) {
        Vertex<V, E> beginV = (Vertex<V, E>) vertexs.get(begin);
        if (beginV == null) return;

        Set<Vertex<V, E>> visitedVertexs = new HashSet<>();
        Stack<Vertex<V, E>> stack = new Stack<>();
        stack.push(beginV);
        visitedVertexs.add(beginV);
        if (visitor.visit(beginV.value)) return;
        while (!stack.isEmpty()) {
            Vertex<V, E> pop = stack.pop();
            for (Edge<V, E> outEdge : pop.outEdges) {
                if (visitedVertexs.contains(outEdge.to)) continue;
                stack.push(outEdge.from);
                stack.push(outEdge.to);
                visitedVertexs.add(outEdge.to);
                if (visitor.visit(outEdge.to.value)) return;
                break;
            }
        }
    }

    @Override
    public List<V> topologicalSort() {
        List<V> list = new ArrayList<>();
        Map<Vertex<V, E>, Integer> ins = new HashMap();
        Queue<Vertex<V, E>> queue = new LinkedList<>();
        vertexs.forEach((V v, Vertex<V, E> vertex) -> {
            int in = vertex.inEdges.size();
            if (in == 0) {
                queue.offer(vertex);
            } else {
                ins.put(vertex, in);
            }
        });

        while (!queue.isEmpty()) {
            Vertex<V, E> poll = queue.poll();
            list.add(poll.value);
            for (Edge<V, E> outEdge : poll.outEdges) {
                int toIns = ins.get(outEdge.to) - 1;
                if (toIns == 0) {
                    queue.offer(outEdge.to);
                } else {
                    ins.put(outEdge.to, toIns);
                }
            }
        }
        return list;
    }

    @Override
    public Set<EdgeInfo<V, E>> mst() {
        return kruskal();
    }

    private Set<EdgeInfo<V, E>> prim() {
        Iterator<Vertex<V, E>> iterator = vertexs.values().iterator();
        if (!iterator.hasNext()) return null;

        Set<EdgeInfo<V, E>> edgeInfos = new HashSet<>();
        Set<Vertex<V, E>> addedVertexs = new HashSet<>();
        Vertex<V, E> begin = vertexs.values().iterator().next();
        addedVertexs.add(begin);

        BinaryHeap<Edge<V, E>> heap = new BinaryHeap<>(begin.outEdges, edgeComparator);
        int vertexsSize = vertexs.size();
        while (!heap.isEmpty() && addedVertexs.size() < vertexsSize) {
            Edge<V, E> edge = heap.remove();
            if (addedVertexs.contains(edge.to)) continue;
            addedVertexs.add(edge.to);
            edgeInfos.add(edge.info());
            heap.addAll(edge.to.outEdges);
        }
        return edgeInfos;
    }

    private Set<EdgeInfo<V, E>> kruskal() {
        int edgeSize = vertexs.size() - 1;
        if (edgeSize == -1) return null;
        Set<EdgeInfo<V, E>> edgeInfos = new HashSet<>();

        BinaryHeap<Edge<V, E>> heap = new BinaryHeap<>(edges, edgeComparator);
        UnionFind<Vertex<V, E>> uf = new UnionFind<>();
        vertexs.forEach((V v, Vertex<V, E> vertex) -> {
            uf.makeSet(vertex);
        });

        while (!heap.isEmpty() && edgeInfos.size() < edgeSize) {
            Edge<V, E> edge = heap.remove();
            if (uf.isSame(edge.from, edge.to)) continue;
            edgeInfos.add(edge.info());
            uf.union(edge.from, edge.to);
        }
        return edgeInfos;
    }

    @Override
    public Map<V, PathInfo<V, E>> shortestPath(V begin) {
        return bellmanFord(begin);
    }

    public Map<V, PathInfo<V, E>> dijkstra(V begin) {
        Vertex<V, E> beginVertex = vertexs.get(begin);
        if (beginVertex == null) return null;
        Map<V, PathInfo<V, E>> selectedPaths = new HashMap<>();
        Map<Vertex<V, E>, PathInfo<V, E>> paths = new HashMap<>();
        for (Edge<V, E> outEdge : beginVertex.outEdges) {
            PathInfo<V, E> path = new PathInfo<>();
            path.setWeight(outEdge.weight);
            path.getEdgeInfos().add(outEdge.info());
            paths.put(outEdge.to, path);
        }
        while (!paths.isEmpty()) {
            Map.Entry<Vertex<V, E>, PathInfo<V, E>> minPathEntry = getMinPath(paths);
            Vertex<V, E> minVertex = minPathEntry.getKey();
            PathInfo<V, E> minPath = minPathEntry.getValue();
            selectedPaths.put(minVertex.value, minPath);
            paths.remove(minVertex);
            for (Edge<V, E> outEdge : minVertex.outEdges) {
                if (selectedPaths.containsKey(outEdge.to.value)) continue;
                relaxForDijkstra(outEdge, minPath, paths);
            }
        }
        selectedPaths.remove(begin);
        return selectedPaths;
    }

    private void relaxForDijkstra(Edge<V, E> edge, PathInfo<V, E> fromPath, Map<Vertex<V, E>, PathInfo<V, E>> paths) {
        PathInfo<V, E> oldPath = paths.get(edge.to);
        E newWeight = weightManager.add(fromPath.getWeight(), edge.weight);
        if (oldPath != null && weightManager.compare(newWeight, oldPath.getWeight()) <= 0) return;
        if (oldPath == null) {
            oldPath = new PathInfo<>();
            paths.put(edge.to, oldPath);
        } else {
            oldPath.getEdgeInfos().clear();
        }
        oldPath.setWeight(newWeight);
        oldPath.getEdgeInfos().addAll(fromPath.getEdgeInfos());
        oldPath.getEdgeInfos().add(edge.info());
    }

    private Map.Entry<Vertex<V, E>, PathInfo<V, E>> getMinPath(Map<Vertex<V, E>, PathInfo<V, E>> paths) {
        Iterator<Map.Entry<Vertex<V, E>, PathInfo<V, E>>> iterator = paths.entrySet().iterator();
        Map.Entry<Vertex<V, E>, PathInfo<V, E>> min = iterator.next();
        while (iterator.hasNext()) {
            Map.Entry<Vertex<V, E>, PathInfo<V, E>> next = iterator.next();
            if (weightManager.compare(next.getValue().getWeight(), min.getValue().getWeight()) > 0) {
                min = next;
            }
        }
        return min;
    }

    private Map<V, PathInfo<V, E>> bellmanFord(V begin) {
        Vertex<V, E> beginVertex = vertexs.get(begin);
        if (beginVertex == null) return null;
        Map<V, PathInfo<V, E>> selectedPaths = new HashMap<>();
        selectedPaths.put(begin, new PathInfo<>(weightManager.zero()));
        int count = vertexs.size() - 1;
        for (int i = 0; i < count; i++) {
            for (Edge<V, E> edge : edges) {
                PathInfo<V, E> fromPath = selectedPaths.get(edge.from.value);
                if (fromPath == null) continue;
                relax(edge, fromPath, selectedPaths);
            }
        }
        for (Edge<V, E> edge : edges) {
            PathInfo<V, E> fromPath = selectedPaths.get(edge.from.value);
            if (fromPath == null) continue;
            if (relax(edge, fromPath, selectedPaths)) {
                System.out.println("有负权环");
                return null;
            }
        }
        selectedPaths.remove(begin);
        return selectedPaths;
    }

    private boolean relax(Edge<V, E> edge, PathInfo<V, E> fromPath, Map<V, PathInfo<V, E>> selectedPaths) {
        PathInfo<V, E> oldPath = selectedPaths.get(edge.to.value);
        E newWeight = weightManager.add(fromPath.getWeight(), edge.weight);
        if (oldPath != null && weightManager.compare(newWeight, oldPath.getWeight()) <= 0) return false;

        if (oldPath == null) {
            oldPath = new PathInfo<>();
            selectedPaths.put(edge.to.value, oldPath);
        } else {
            oldPath.getEdgeInfos().clear();
        }

        oldPath.setWeight(newWeight);
        oldPath.getEdgeInfos().addAll(fromPath.getEdgeInfos());
        oldPath.getEdgeInfos().add(edge.info());
        return true;
    }

    @Override
    public Map<V, Map<V, PathInfo<V, E>>> shortestPath() {
        Map<V, Map<V, PathInfo<V, E>>> paths = new HashMap<>();
        for (Edge<V, E> edge : edges) {
            Map<V, PathInfo<V, E>> map = paths.get(edge.from.value);
            if (map == null) {
                map = new HashMap<>();
                paths.put(edge.from.value, map);
            }
            PathInfo<V, E> pathInfo = new PathInfo<>();
            pathInfo.setWeight(edge.weight);
            pathInfo.getEdgeInfos().add(edge.info());
            map.put(edge.to.value, pathInfo);
        }
        vertexs.forEach((V v2, Vertex<V, E> vertex2) -> {
            vertexs.forEach((V v1, Vertex<V, E> vertex1) -> {
                vertexs.forEach((V v3, Vertex<V, E> vertex3) -> {
                    if (v1.equals(v2) || v1.equals(v3) || v2.equals(v3)) return;
                    PathInfo<V, E> path12 = getPathInfo(v1, v2, paths);
                    if (path12 == null) return;
                    PathInfo<V, E> path23 = getPathInfo(v2, v3, paths);
                    if (path23 == null) return;
                    PathInfo<V, E> path13 = getPathInfo(v1, v3, paths);
                    E newWeight = weightManager.add(path12.getWeight(), path23.getWeight());
                    if (path13 != null && weightManager.compare(newWeight, path13.getWeight()) <= 0) return;
                    if (path13 == null) {
                        path13 = new PathInfo<>();
                        paths.get(v1).put(v3, path13);
                    } else {
                        path13.getEdgeInfos().clear();
                    }
                    path13.setWeight(newWeight);
                    path13.getEdgeInfos().addAll(path12.getEdgeInfos());
                    path13.getEdgeInfos().addAll(path23.getEdgeInfos());
                });
            });
        });
        return paths;
    }

    private PathInfo<V, E> getPathInfo(V from, V to, Map<V, Map<V, PathInfo<V, E>>> paths) {
        Map<V, PathInfo<V, E>> map = paths.get(from);
        return map == null ? null : map.get(to);
    }

    @Override
    public void print() {
        System.out.println("[顶点]-------------------");
        vertexs.forEach((V v, Vertex<V, E> vertex) -> {
            System.out.println(v);
            System.out.println("out-----------");
            System.out.println(vertex.outEdges);
            System.out.println("in-----------");
            System.out.println(vertex.inEdges);
        });

        System.out.println("[边]-------------------");
        edges.forEach((Edge<V, E> edge) -> {
            System.out.println(edge);
        });
    }

    private static class Vertex<V, E> {
        V value;
        Set<Edge<V, E>> inEdges = new HashSet<>();
        Set<Edge<V, E>> outEdges = new HashSet<>();

        Vertex(V value) {
            this.value = value;
        }

        @Override
        public boolean equals(Object obj) {
            return Objects.equals(value, ((Vertex<V, E>) obj).value);
        }

        @Override
        public int hashCode() {
            return value == null ? 0 : value.hashCode();
        }

        @Override
        public String toString() {
            return this.value == null ? "null" : value.toString();
        }
    }

    private static class Edge<V, E> {
        Vertex<V, E> from;
        Vertex<V, E> to;
        E weight;

        Edge(Vertex<V, E> from, Vertex<V, E> to) {
            this.from = from;
            this.to = to;
        }

        EdgeInfo<V, E> info() {
            return new EdgeInfo<>(from.value, to.value, weight);
        }

        @Override
        public boolean equals(Object obj) {
            Edge<V, E> edge = (Edge<V, E>) obj;
            return Objects.equals(from, edge.from) && Objects.equals(to, edge.to);
        }

        @Override
        public int hashCode() {
            return from.hashCode() * 31 + to.hashCode();
        }

        @Override
        public String toString() {
            return from + "->" + to + ":" + weight;
        }
    }
}
